#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import json
from typing import Dict, List, Any

class AverageMetricsUpdater:
    """
    更新现有实验结果文件中的平均指标，包括嵌套对象中的指标
    """
    
    def __init__(self, results_dir: str):
        """
        初始化更新器
        
        Args:
            results_dir: 实验结果文件夹路径
        """
        self.results_dir = results_dir
    
    def collect_metrics(self, data: Any, metrics_accumulator: Dict[str, float], 
                       metrics_count: Dict[str, int], prefix: str = ''):
        """
        递归收集嵌套对象中的数值指标
        
        Args:
            data: 要收集指标的数据
            metrics_accumulator: 指标累加器字典
            metrics_count: 指标计数字典
            prefix: 当前路径前缀
        """
        if isinstance(data, dict):
            for key, value in data.items():
                new_prefix = f"{prefix}{key}_" if prefix else key
                self.collect_metrics(value, metrics_accumulator, metrics_count, new_prefix)
        elif isinstance(data, (int, float)):
            if not prefix.startswith('__'):
                if prefix not in metrics_accumulator:
                    metrics_accumulator[prefix] = 0.0
                    metrics_count[prefix] = 0
                metrics_accumulator[prefix] += data
                metrics_count[prefix] += 1
    
    def calculate_average_metrics(self, results: List[Dict[str, Any]]) -> Dict[str, float]:
        """
        计算实验结果中所有数值指标的平均值，包括嵌套对象中的指标
        
        Args:
            results: 实验结果列表
        
        Returns:
            各个指标的平均值字典
        """
        if not results:
            return {}
        
        # 收集所有数值类型的指标
        metrics_accumulator = {}
        metrics_count = {}
        
        for result in results:
            # 跳过错误记录
            if 'error' in result:
                continue
            
            # 递归收集所有数值指标
            self.collect_metrics(result, metrics_accumulator, metrics_count)
        
        # 计算平均值
        average_metrics = {}
        for key, total in metrics_accumulator.items():
            if metrics_count[key] > 0:
                average_metrics[key] = round(total / metrics_count[key], 4)
        
        return average_metrics
    
    def update_all_results_files(self):
        """
        更新所有结果文件的平均指标
        """
        # 遍历结果文件夹中的所有文件
        for filename in os.listdir(self.results_dir):
            if filename.startswith('experiment_results_window_') and filename.endswith('.json'):
                file_path = os.path.join(self.results_dir, filename)
                try:
                    # 读取文件内容
                    with open(file_path, 'r', encoding='utf-8') as f:
                        data = json.load(f)
                        
                    if data and isinstance(data, list) and len(data) > 0:
                        # 获取第一个实验组（应该只有一个）
                        experiment_data = data[0]
                        
                        if 'results' in experiment_data:
                            # 计算新的平均指标
                            average_metrics = self.calculate_average_metrics(experiment_data['results'])
                            
                            # 更新实验数据中的平均指标
                            experiment_data['average_metrics'] = average_metrics
                            
                            # 写回文件
                            with open(file_path, 'w', encoding='utf-8') as f:
                                json.dump(data, f, ensure_ascii=False, indent=2)
                            
                            print(f"已更新文件 {filename} 的平均指标")
                            print(f"  包含的指标数量: {len(average_metrics)}")
                            
                            # 打印一些重要的指标示例
                            important_metrics = ['text_metrics_cosine_similarity', 'text_metrics_bertscore', 
                                               'text_metrics_bleu_bleu-1']
                            for metric in important_metrics:
                                if metric in average_metrics:
                                    print(f"  {metric}: {average_metrics[metric]:.4f}")
                except Exception as e:
                    print(f"处理文件 {filename} 失败: {str(e)}")


def main():
    """主函数"""
    # 创建更新器实例
    updater = AverageMetricsUpdater('../results')
    
    # 更新所有结果文件的平均指标
    updater.update_all_results_files()
    
    print("\n所有文件更新完成！您现在可以使用 analyze_results.py 脚本来查看详细的指标比较。")


if __name__ == '__main__':
    main()